//||=========================================================================||
//||                                                                         ||
//||                             MolDynSys.cpp                               ||
//||                                                                         ||
//||                           by the_best_goose))                           ||
//||                                                                         ||
//||=========================================================================||


#include <math.h>
#include <assert.h>
#include <cstdlib>
#include <algorithm>
#include <time.h>
#include "MolDynSys.h"

using std::rand;
using std::srand;

//-----------------------------------------------------------------------------

MolDynSys::MolDynSys(unsigned int part_count, float density, Tests test):
    coords_(),
    vels_(),
    forces_(),
    part_count_(part_count),
    size_(),
    energy_kin_(0),
    energy_pot_(0),
    thermostat_status_(false),
    need_temperature_(0),
    current_temperature_(0),
    temperature_temp_(0),
    last_temperature_update_(0),
    temperature_up_time_(1000),
    plog_(NULL)

{
    plog_ = fopen("graph/log", "w");
    assert(plog_);
    if(test != NO_TEST)
        Test(test);
    assert(part_count_ > 0);
    assert(density > 0);

    size_[0] = size_[1] = size_[2] = powf(part_count_ / density, 1./3);    
    Init(part_count_);

    fprintf(plog_, "#K(t)\t #U(t)\n");

}

//-----------------------------------------------------------------------------

MolDynSys::~MolDynSys()
{
    Destroy();
    assert(plog_);
    fclose(plog_);
}

//-----------------------------------------------------------------------------

void MolDynSys::Init(int part_count)
{
    part_count_ = part_count;
    for(int i = 0; i < 3; ++i)
    {
        coords_[i] = new float[part_count_];
        vels_[i] = new float[part_count_];
        forces_[i] = new float[part_count_];
        for(unsigned int j = 0; j < part_count_; ++j)
            coords_[i][j] = vels_[i][j] = forces_[i][j] = 0; 
    }

}

//-----------------------------------------------------------------------------

void MolDynSys::Destroy()
{
    for(int i = 0; i < 3; ++i)
    {
        assert(coords_[i]);
        assert(vels_[i]);
        assert(forces_[i]);

        delete[] (coords_[i]);
        delete[] (vels_[i]);
        delete[] (forces_[i]);
    }
}

//-----------------------------------------------------------------------------

void MolDynSys::InitSystem()
{
    time_t t;
    srand(time(&t));
    int num_per_side = pow(part_count_, 1./3) + 1;
    float ip = size_[0] / num_per_side;
    float jp = size_[1] / num_per_side;
    float kp = size_[2] / num_per_side;

    unsigned int n = 0;
    for(float i = 0; i < size_[0]; i += ip)
        for(float j = 0; j < size_[1]; j += jp) 
            for(float k = 0; k < size_[2]; k += kp)
            {
                if(n >= part_count_)
                    break;
                coords_[0][n] = i;
                coords_[1][n] = j;
                coords_[2][n] = k;
                n++;
            }

    energy_kin_ = CountKinEnergy();

}

//-----------------------------------------------------------------------------

void MolDynSys::NextStep(float t)
{
    for(int i = 0; i < 3; ++i)
        for(unsigned int j = 0; j < part_count_; ++j)
            if(coords_[i][j] < 0)
                coords_[i][j] += size_[i];
            else if(coords_[i][j] > size_[i])
                coords_[i][j] -= size_[i];

    CountForces();
    for(int i = 0; i < 3; ++i)
        for(unsigned int j = 0; j < part_count_; ++j)
        {
            vels_[i][j] += forces_[i][j] * t;
            coords_[i][j] += vels_[i][j] * t;
        }

    energy_kin_ = CountKinEnergy();
    if(thermostat_status_)
        UseThermostat(t);
}

//-----------------------------------------------------------------------------

float MolDynSys::CountKinEnergy()
{
    float energy = 0;
    for(int i = 0; i < 3; ++i)
        for(unsigned int j = 0; j < part_count_; ++j)
            energy += vels_[i][j]*vels_[i][j] / 2;
   
    temperature_temp_ += energy;
    last_temperature_update_++;

    if(last_temperature_update_ >= temperature_up_time_)
    {
        current_temperature_ =  temperature_temp_ /
                                (last_temperature_update_ * 1.5 * part_count_);
        temperature_temp_ = last_temperature_update_ = 0;
    }

    return energy;
         
}

//-----------------------------------------------------------------------------

float MolDynSys::GetKinEnergy()
{
    return energy_kin_;
}

//-----------------------------------------------------------------------------

float MolDynSys::GetPotEnergy()
{
    return energy_pot_;
}

//-----------------------------------------------------------------------------

float MolDynSys::GetEnergy()
{
    return energy_kin_ + energy_pot_;
}

//-----------------------------------------------------------------------------

void MolDynSys::Test(Tests test_number)
{
    if(test_number == ONE_PARTICLE)
    {
        Init(1);

        size_[0] = size_[1] = size_[2] = 10.0f;
        coords_[0][0] = coords_[1][0] = coords_[2][0] = 0.0f;
        vels_[0][0] = 1.0f;
        vels_[0][1] = vels_[0][2] = 0.0f;

        FILE* fout = fopen("OneParticleTest.out", "w");
        assert(fout);

        for(float i = 1; i < 100; i += 0.1)
        {
            NextStep(0.1);
            fprintf(fout, "%f ", coords_[0][0]);
        }
        fclose(fout); 
        Destroy();
   
    }
    else if(test_number == TWO_PARTICLES)
    {
        float dx = -2.8;
        Init(2);

        size_[0] = size_[1] = size_[2] = 100.0f;
        coords_[0][0] = 0.0f + dx;
        coords_[0][1] = -1.8f + dx;
        vels_[0][0] = 0.0f;
        vels_[0][1] = 0.0f;
        
        FILE* dr_fout = fopen("TwoParticleTest_coords_delta.out", "w");
        FILE* en_pot_fout = fopen("TwoParticleTest_energy_pot.out", "w");
        FILE* en_kin_fout = fopen("TwoParticleTest_energy_kin.out", "w");
        assert(dr_fout);
        assert(en_pot_fout);
        assert(en_kin_fout);

        for(int i = 1; i < 10000; ++i)
        {
            NextStep(0.001);
 
            energy_kin_ = CountKinEnergy();

            printf("%f\t%f\t%f\n", energy_kin_, energy_pot_, sqrt(Dist2(0,1)));
            fprintf(dr_fout, "%f ", coords_[0][1] - coords_[0][0]);
            fprintf(en_pot_fout, "%f ", energy_pot_); 
            fprintf(en_kin_fout, "%f ", energy_kin_);
        }


        fclose(dr_fout);
        fclose(en_pot_fout);
        fclose(en_kin_fout);
        Destroy();
    }
    else if(test_number == THREE_PARTICLES)
    {
        Init(3);
        float dx = 0.2f;

        size_[0] = size_[1] = size_[2] = 100.0f;
        coords_[0][0] = -1.8f + dx;
        coords_[0][1] = -0.0f + dx;
        coords_[0][2] = 1.8f + dx;

        for(int i = 1; i < 40000; ++i)
        {
            NextStep(0.01);

            energy_kin_ = CountKinEnergy();

            printf("%f\t%f\t%f\t\n", energy_kin_, energy_pot_,
                                       energy_kin_ + energy_pot_ );
        }

        Destroy();
    }

    else if(test_number == MANY_PARTICLES)
    {
        Init(8);

        size_[0] = size_[1] = size_[2] = 15.0f;
        for(unsigned int i = 0; i < part_count_; ++i)
            coords_[0][i] = i * size_[0] /(part_count_ + 1.0f);
        
        for(int i = 1; i < 500000; ++i)
        {
            NextStep(0.01);

            printf("%f\t%f\t%f\n", energy_kin_, energy_pot_,
                                       energy_kin_ + energy_pot_);
        }


        Destroy();
    }
}

//-----------------------------------------------------------------------------

inline void MolDynSys::CountForces()
{
    float dist2 = 0, force = 0, force_pr = 0, sr12 = 0, sr6 = 0,
          sr2 = 0;
    float s2 = sigma * sigma;
    float temp = 0;
    NullArray(forces_); 
    
    energy_pot_ = 0;
    for(unsigned int i = 0; i < part_count_; ++i)
        for(unsigned int j = i + 1; j < part_count_; ++j)
        {
            dist2 = Dist2(i, j);
            if(dist2 >= r_cr2_)
                continue;

            sr2 = s2 / dist2; 
            sr6 = sr2 * sr2 * sr2;
            sr12 = sr6 * sr6;
            
            energy_pot_ += 4 * eps * (sr12 - sr6);
            force = 24 * eps / sigma * sr2 * (2 * sr12 - sr6);

            for(int k = 0; k < 3; ++k)
            {
                temp = coords_[k][j] - coords_[k][i];
                if(fabs(temp) < size_[k] - fabs(temp))
                    force_pr = force * temp;
                else{ if(temp >= 0)
                    force_pr = force * (temp - size_[k]);
                else if(temp < 0)
                    force_pr = force * (temp + size_[k]);
                }    forces_[k][i] -= force_pr;
                forces_[k][j] += force_pr;
            }
        }
 
}

//-----------------------------------------------------------------------------

inline float MolDynSys::Dist2(int i, int j)
{
    float ret = 0;
    float temp_r = 0;
    float delta_r = 0;
    for(int a = 0; a < 3; ++a)
    {
        delta_r = fabs(coords_[a][i] - coords_[a][j]);
        if(delta_r < fabs(size_[a] - delta_r))
            temp_r = delta_r;
        else
            temp_r = size_[a] - delta_r;

        ret += temp_r * temp_r;
    }

    assert(ret >= 0.4);
    return ret;
}

//-----------------------------------------------------------------------------

inline void MolDynSys::NullArray(float** arr)
{
    for(int i = 0; i < 3; ++i)
        for(unsigned int j = 0; j < part_count_; ++j)
            arr[i][j] = 0.0f;
}

//-----------------------------------------------------------------------------

void MolDynSys::Print()
{
    fprintf(plog_, "%.8E\t %.8E\t %.8E\n", energy_kin_ / part_count_,
                                           energy_pot_ / part_count_,
                                           (energy_kin_ + energy_pot_) /
                                           part_count_);
}

//-----------------------------------------------------------------------------

void MolDynSys::RunThermostat(float T, float eps, float step)
{
    need_temperature_ = T;
    eps *= part_count_;
    T *= part_count_;

    EnableThermostat();
    while(fabs(energy_kin_ -  1.5 * T) >= eps)
    {
        NextStep(step);
        Print();
        printf("Thermostating... delta_r = %f\r", current_temperature_ - T);
    }
    DisableThermostat();
}

//-----------------------------------------------------------------------------

void MolDynSys::EnableThermostat()
{
    thermostat_status_ = true;
}

//-----------------------------------------------------------------------------

void MolDynSys::DisableThermostat()
{
    thermostat_status_ = false;
}

//-----------------------------------------------------------------------------

void MolDynSys::SetTemperatureUpTime(float up_time)
{
    temperature_up_time_ = up_time;
}

//-----------------------------------------------------------------------------

float MolDynSys::GetTemperature()
{
    return current_temperature_;
}

//-----------------------------------------------------------------------------

void MolDynSys::UseThermostat(float t)
{
    for(int j = 0; j < 3; ++j)
        for(unsigned int i = 0; i < part_count_; ++i)
            vels_[j][i] *= sqrt(1 + t * 0.0005 *  
                         (need_temperature_ * 1.5 * part_count_ - energy_kin_));
}

//-----------------------------------------------------------------------------

void MolDynSys::Relax(float temperature_eps, float step)
{
    int i = 0;
    float current_temperature_old = current_temperature_;
    for(int i = 0; i < 5; ++i)
        do
        {
            current_temperature_old = current_temperature_;
            do
            {
            printf("Relaxing... \r");
            NextStep(step);
            Print();
            }
            while(last_temperature_update_ < temperature_up_time_- 1);
        }
        while(fabs(current_temperature_old - current_temperature_) >= temperature_eps );
}


